import json
import os
import re
import time
import traceback

import util.common
from util.llm import http_request

MODEL_NAME = os.getenv('MODEL_NAME', "deepseek-chat")

SYSTEM_PROMPT = """你是一名企业安全专家，专门处理企业问题。以下我会给出一段系统问题描述，请帮我提取出以下实体。如果有，就返回它们的序号。结果以json array形式提供
1. ip地址
2. 用户名
3. url
4. 明文密码、弱密码、使用url传递密码、接口返回内容中包含密码
5. 配置文件、代码相关文件、git文件、svn文件
6. 目录浏览
以上为目前的6种实体，如果在描述中出现了上面的或相关实体，请写出它的序号，可能包含多种。
结果示例：
输入：检测到零信任高频访问，统计结果共命中449条日志： [访问次数：5647次，用户名：****@asi**o.com，请求url：192.168.12.2:222]，触发条件'访问次数' > 1500
输出： [1,2]
以下是用户问题：
"""

PROMPT = """以下我将给出一个问题描述，请帮我提取关键词用于搜索引擎搜索。结果用json array表示
问题：{}"""


async def query_rewrite(query):
    a = time.time()
    result = await http_request(system_prompt=SYSTEM_PROMPT, history=[], prompt=PROMPT.format(query),
                                # model="hunyuan-turbo")
                                model=MODEL_NAME)
    b = time.time()
    print("rewrite cost: {}".format(b - a))
    try:
        result = util.common.remove_markdown_code_block(result)
        result = json.loads(result)
        result.append(query)
        result = list(set(result))
    except Exception as e:
        print(traceback.format_exc())
        print("query_rewrite error:{}".format(e))
        result = [query, ]
    print("rewrite result:{}".format(result))
    return result
